# 安装：pip install xgboost
from types import TracebackType
from xgboost import XGBClassifier
from xgboost import XGBRegressor
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from xgboost import plot_tree
import matplotlib.pyplot as plt
from xgboost import to_graphviz  # 将树导出到dot文件

# 设置字体
plt.rcParams["font.sans-serif"] = "SimHei"
# 默认可以显示负号，增加字体显示后。需对负号正常显示进行设置
plt.rcParams['axes.unicode_minus'] = False

# 使用鸢尾花做分类

iris = load_iris()

X, y = iris.data, iris.target

print("特征\n", X)
print("特征名称", iris.feature_names)

X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.2,
    random_state=1,
    stratify=y
)

xgb = XGBClassifier(
    # 一、通用参数
    # booster='gbtree'  # 设置基础模型是谁？可以树 可以线性模型
    # 二、booster参数设置，比如树相关的深度、叶子节点
    # 三、任务参数， 主要设置目标函数
    eval_metric='mlogloss',
    use_label_encoder=False,
)

xgb.fit(X_train, y_train)

print(xgb.score(X_test, y_test))


# 参数1：booster 集成算法对象
# num_trees 第几棵树
# plot_tree(xgb, num_trees=1, fmap='iris_xgb.fmap')
# plt.show()

dot_file = to_graphviz(xgb, num_trees=1, fmap='iris_xgb_v2.fmap')
dot_file.render("iris_xgb.dot")


